from typing import Callable, Iterable
import openai
import traceback
from openai.types.chat.chat_completion_message_tool_call import Function
from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessage, ChatCompletionMessageToolCall
from openai.types.chat.chat_completion import Choice
from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk, ChoiceDeltaToolCall

def retry_api(completion_create: Callable, **kwargs):
    error = None
    for _ in range(3):
        try:
            return completion_create(**kwargs)
        except openai.APIError as e:
            if error is None:
                error = e.message
            traceback.print_exc()
    raise RuntimeError('OpenAI API Error: {}'.format(error))

client = openai.OpenAI()

def chat_completion_create(**kwargs) -> ChatCompletion:
    completion: ChatCompletion = retry_api(client.chat.completions.create, **kwargs)
    # if completion.usage:
    #     print('Usage: {}+{} tokens ({:.2f}x) @{}'.format(completion.usage.prompt_tokens, completion.usage.completion_tokens, completion.usage.prompt_tokens / max(1, completion.usage.completion_tokens), kwargs.get('model', 'gpt-3.5-turbo')))
    return completion

def chat_completion_create_streamed(callback: Callable[[ChatCompletion], None], **kwargs) -> ChatCompletion:
    completion = ChatCompletion(id='', choices=[], created=0, model='', object='chat.completion')
    completions: Iterable[ChatCompletionChunk] = retry_api(client.chat.completions.create, **kwargs, stream=True)
    choices: list[Choice] = []
    for chunk in completions:
        if not choices:
            choices = [Choice(finish_reason='length', index=i, message=ChatCompletionMessage(role='assistant')) for i in range(len(chunk.choices))]
        for i, c in enumerate(chunk.choices):
            assert isinstance(c, ChoiceChunk)
            if c.delta.content:
                con = choices[i].message.content
                if con is None:
                    choices[i].message.content = c.delta.content
                else:
                    choices[i].message.content = con + c.delta.content
            if c.delta.tool_calls:
                if choices[i].message.tool_calls is None:
                    tcs = []
                    for tc in c.delta.tool_calls:
                        assert isinstance(tc, ChoiceDeltaToolCall)
                        assert tc.function
                        mtc = ChatCompletionMessageToolCall(id=tc.id or '', function=Function(name=tc.function.name or '', arguments=tc.function.arguments or ''), type='function')
                        tcs.append(mtc)
                    choices[i].message.tool_calls = tcs
                else:
                    tcs = choices[i].message.tool_calls
                    assert tcs
                    idset = {tc.id: t for t, tc in enumerate(tcs)}
                    for tc in c.delta.tool_calls:
                        assert isinstance(tc, ChoiceDeltaToolCall)
                        tcid = tc.id
                        if tcid and tcid not in idset:
                            assert tc.function
                            mtc = ChatCompletionMessageToolCall(id=tcid, function=Function(name=tc.function.name or '', arguments=tc.function.arguments or ''), type='function')
                            tcs.append(mtc)
                    for tc in c.delta.tool_calls:
                        assert isinstance(tc, ChoiceDeltaToolCall)
                        tcid = tc.id
                        if not tcid or tcid in idset:
                            t = idset[tcid] if tcid else -1
                            mtc = tcs[t]
                            if tc.id:
                                mtc.id = tc.id
                            if tc.function:
                                if tc.function.name:
                                    mtc.function.name += tc.function.name
                                if tc.function.arguments:
                                    mtc.function.arguments += tc.function.arguments
                            if tcid:
                                idset[tcid] = t
                    choices[i].message.tool_calls = tcs
            if c.finish_reason:
                choices[i].finish_reason = c.finish_reason
        completion = ChatCompletion(id=chunk.id, choices=choices, created=chunk.created, model=chunk.model, object='chat.completion')
        callback(completion)
    return completion

if __name__ == "__main__":
    res = chat_completion_create_streamed(print, model='gpt-3.5-turbo', messages=[{'role': 'user', 'content': 'Weather in Beijing and Shenzhen?'}], tools=[
        {"type": "function", "function": {
            "name": "get_weather",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {"type": "string"},
                },
                "required": ["location"],
            },
        }}
    ])
    print(res)
